# -*- encoding: utf-8 -*-
import os
import time

import cv2


def image_extractor(video_path, jump_fps=1, count=None, pic_size=(1920, 1080)):
    """
    :param video_path: 视频路径
    :param jump_fps: 跳帧数
    :param count: 截取图片数量
    :param pic_size: 截取图片大小
    :return:
    """
    cap = cv2.VideoCapture(video_path)
    print("PROP_FPS: ", round(cap.get(cv2.CAP_PROP_FPS)))
    print("PROP_FRAME_COUNT: ", cap.get(cv2.CAP_PROP_FRAME_COUNT))
    f_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    f_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    frame_num = 0
    frame_cut = 0
    full_path, full_name = os.path.split(video_path)
    name, _ = os.path.splitext(full_name)
    storage_path = os.path.join(full_path, "ImagesComplete")
    if not os.path.exists(storage_path):
        os.makedirs(storage_path)
    if cap.isOpened():
        if count is not None:
            while frame_cut < count:
                _, frame = cap.read()
                if frame is None:
                    break
                if f_width != pic_size[0] and f_height != pic_size[1]:
                    frame = cv2.resize(frame, pic_size)
                frame_num += 1
                if frame_num % jump_fps == 0:
                    frame_cut += 1
                    print(f"number {frame_num} frame")
                    print(f'write {frame_cut} frame')
                    file_full_path = os.path.join(storage_path,
                                                  name + "_" + str(frame_cut) + ".jpg")
                    print(file_full_path)
                    cv2.imwrite(file_full_path, frame)
        else:
            while True:
                _, frame = cap.read()
                if frame is None:
                    break
                if f_width != pic_size[0] and f_height != pic_size[1]:
                    frame = cv2.resize(frame, pic_size)
                if frame_num % jump_fps == 0:
                    frame_cut += 1
                    print(f"number {frame_num + 1} frame")
                    print(f'write {frame_cut} frame')
                    file_full_path = os.path.join(storage_path,
                                                  name + "_" + str(frame_cut) + ".jpg")
                    print(file_full_path)
                    cv2.imwrite(file_full_path, frame)
                frame_num += 1
    else:
        print("Open Video Failed!")
    cap.release()


def image_extractor_by_sec(video_path, start_point, end_point, jump_sec=None, pic_size=(1920, 1080)):
    cap = cv2.VideoCapture(video_path)
    if cap.isOpened():
        f_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        f_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        fps = round(cap.get(cv2.CAP_PROP_FPS))
        frame_number = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        duration = frame_number / fps
        start_frame = int(fps * start_point)
        end_frame = int(fps * end_point)
        jump_frame = int(fps * jump_sec)
        frame_count = 0
        full_path, full_name = os.path.split(video_path)
        name, _ = os.path.splitext(full_name)
        storage_path = os.path.join(full_path, "ImagesComplete")
        if not os.path.exists(storage_path):
            os.makedirs(storage_path)
        while True:
            _, frame = cap.read()
            if frame is None:
                break
            # 判断图片大小是否符合期望值
            if f_width != pic_size[0] and f_height != pic_size[1]:
                frame = cv2.resize(frame, pic_size)
            # 判断帧数值是否在期望区间
            if frame_count < int(start_frame):
                print('Counting number ' + str(frame_count) + ' frequency')
            elif int(start_frame) <= frame_count <= int(end_frame):
                # 判断跳帧值是否符合期望数值
                if (frame_count - start_frame) % jump_frame == 0:
                    print('Writing number ' + str(frame_count) + ' frequency')
                    file_full_path = os.path.join(storage_path,
                                                  name + "_" + str(frame_count) + ".jpg")
                    cv2.imwrite(file_full_path, frame)
            else:
                break
            frame_count += 1
        cap.release()
        print("FPS: ", fps)
        print("f_shape: ", f_width, f_height)
        print(str(start_point // 60) + ":" + str(start_point % 60) + " - " + str(end_point // 60) + ":" + str(
            end_point % 60))
        print("Frame Number: ", frame_number)
        print("Duration: ", duration)
    else:
        print('video opened failed !')


if __name__ == "__main__":
    time_01 = time.time()
    video_store_path = r"/home/kaijia/algo-env/Utils/sample/bridge_on_off"
    video_store_list = os.listdir(video_store_path)
    # video_store_list = ["output_2024-03-20_10-22-35.mp4"]
    for tmp_video in video_store_list:
        full_video_path = os.path.join(video_store_path, tmp_video)
        tmp_interval = 25
        pic_size = (1920, 1080)
        image_extractor(full_video_path, tmp_interval, pic_size=pic_size)
        # start_time = 10 * 60 + 40
        # end_time = 11 * 60 + 20
        # jump_sec = 5
        # pic_size = (1920, 1080)
        # image_extractor_by_sec(full_video_path,
        #                        start_time, end_time,
        #                        jump_sec,
        #                        pic_size=pic_size)
    time_02 = time.time()
    print("total time ", time_02 - time_01, " s")
